package com.xiashitech.agent.instrumentation.base;

import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.AttributesExtractor;
import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.ContextCustomizer;
import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.SpanNameExtractor;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.Attributes;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;

import javax.annotation.Nullable;
import java.util.function.Consumer;
import java.util.logging.Logger;

/**
 * 插桩业务逻辑处理器
 * @author zhaoyk
 * @param <REQUEST>
 * @param <RESPONSE>
 */
abstract public class InstProcessor<REQUEST, RESPONSE> implements AttributesExtractor<REQUEST, RESPONSE>, ContextCustomizer<REQUEST>, SpanNameExtractor<REQUEST> {

    protected static final Logger logger = Logger.getLogger("SpanAttributesExtractor");

    protected String spanName;

    /**
    public InstProcessor() {
        this.spanName = getProcessorName();
    }
    **/

    public InstProcessor(String spanName) {
        this.spanName = spanName;
    }

    public String getProcessorName() {
        return getClass().getName();
    }

    @Override
    public void onStart(AttributesBuilder attributesBuilder, Context context, REQUEST request) {
        DebugUtils.log(logger, getProcessorName() + " - start - " + context.toString());
        onXSStart(attributesBuilder, context, request);
    }

    /**
     * 业务逻辑的onStart处理，可以执行获取baggage信息，设置attribute等操作
     */
    public abstract void onXSStart(AttributesBuilder attributesBuilder, Context context, REQUEST request);

    @Override
    public Context onStart(Context context, REQUEST request, Attributes attributes) {
        Context result = setBaggageAndContext(context, request, attributes);
        return result != null ? result : context;
    }

    /**
     * 在start阶段会调用该方法，更新baggage信息的操作应放在此处执行
     * @return 新的context对象，如果返回null则不更新context
     */
    public abstract Context setBaggageAndContext(Context context, REQUEST request, Attributes attributes);

    /**
     * 公用方法，更新baggage数据，并返回新的context
     */
    protected Context updateBaggage(Context context, Consumer<BaggageBuilder> updateAction) {
        BaggageBuilder baggageBuilder = Baggage.fromContext(context).toBuilder();
        updateAction.accept(baggageBuilder);
        return baggageBuilder.build().storeInContext(context);
    }

    @Override
    public void onEnd(AttributesBuilder attributesBuilder, Context context, REQUEST request, @Nullable RESPONSE response, @Nullable Throwable throwable) {
        onXSEnd(attributesBuilder, context, request, response, throwable);
        DebugUtils.log(logger, getProcessorName() + " - end - " + context.toString());
    }

    /**
     * 业务逻辑的onXSEnd处理，可执行设置TraceContextVars的操作，以实现向后续节点传递数据
     */
    public abstract void onXSEnd(AttributesBuilder attributesBuilder, Context context, REQUEST request, RESPONSE response, Throwable throwable);

    @Override
    public String extract(REQUEST request) {
        return spanName;
    }

    public Instrumenter<REQUEST, RESPONSE> buildInstrumenter(String instName) {
        return Instrumenter.<REQUEST, RESPONSE>builder(GlobalOpenTelemetry.get(), instName, this)
                .addAttributesExtractor(this).addContextCustomizer(this)
                .buildInstrumenter();
    }

    public static <REQUEST, RESPONSE> InstProcessor<REQUEST, RESPONSE> NoopProcessor(String spanName) {

        return new InstProcessor<REQUEST, RESPONSE>(spanName) {

            @Override
            public void onXSStart(AttributesBuilder attributesBuilder, Context context, REQUEST request) {
            }

            @Override
            public Context setBaggageAndContext(Context context, REQUEST request, Attributes attributes) {
                return null;
            }

            @Override
            public void onXSEnd(AttributesBuilder attributesBuilder, Context context, REQUEST request, RESPONSE response, Throwable throwable) {
            }

        };
    }

}
